#!/usr/bin/env drgn

from drgn.helpers.linux.cgroup import css_for_each_child,cgroup_path
from drgn import Object
from drgn.helpers.linux.pid import find_task
from drgn.helpers.linux.list import list_for_each_entry,list_for_each
from drgn import container_of
from drgn.helpers.linux.percpu import per_cpu

import sys
import argparse

def print_cgroup_from_addr(addr):
    parent = int(addr, 16)
    print(f'parent: {parent:#x}')

    pcss = Object(prog, 'struct cgroup_subsys_state', address=parent)

    for css in css_for_each_child(pcss.address_of_()):
        print(f'flags: {css.flags.value_():#x}')
        #print(f'cgroup: {css.cgroup.value_():#x}')
        #name = cgroup_name(css.cgroup)
        #print(f'{name}')
        path = cgroup_path(css.cgroup).decode('utf-8')
        print(f'path: {path}')

def print_cgroup_from_pid(pid):
    task = find_task(prog, pid=pid)

    print_cgroup_task(task)

def print_cgroup_task(task):

    comm = task.comm.string_().decode('utf-8')
    pid = task.pid.value_()
    ppid = task.parent.pid.value_()
    print(f'task: {hex(task)} pid: {pid} comm: {comm} ppid: {ppid}')


    cgroup_roots = prog['cgroup_roots'].address_of_()
    dfl_root = prog['cgrp_dfl_root'].address_of_().value_()
    global_css_set = prog['cgroup_subsys']
    ssid_end = prog['CGROUP_SUBSYS_COUNT']

    print('{:<20}{}'.format('Cgroup Name', 'Cgroup Path'))
    for pos in list_for_each(cgroup_roots):

        cg_root_type = prog.type('struct cgroup_root')
        offset = cg_root_type.member('root_list').offset
        #print(f'offset: {offset:#x}')
        root_addr = pos.value_() - offset

        if root_addr == dfl_root:
            continue

        root = container_of(pos, 'struct cgroup_root', 'root_list')
        count = 0
        #print(f'root_addr: {root_addr:#x}')
        cg_name = ''
        for ssid in range(0, ssid_end):
            if (root.subsys_mask & (1 << ssid)) != 0:
                css = global_css_set[ssid].read_()
                cg_name = '{}{}{}'.format(cg_name, ',' if count else '', css.legacy_name.string_().decode('utf-8'))
                count += 1

        if cg_name != '':
            print(f'{cg_name:<20}', end='')

        if root.name.string_().decode('utf-8') != '':
            cg_name = '{}name={}'.format(',' if count else '', root.name.string_().decode('utf-8'))
            print(f'{cg_name:<20}', end='')

        cgroup = None
        css_set = task.cgroups.read_()
        for link in list_for_each_entry('struct cgrp_cset_link', css_set.cgrp_links.address_of_(), 'cgrp_link'):
            cg = link.cgrp.read_()
            cg_root_addr = cg.root.value_()
            #print(f'cg_root_addr: {cg_root_addr:#x}')
            if cg_root_addr == root_addr:
                cgroup = cg

        path = cgroup_path(cgroup).decode('utf-8')
        print(f'{path}')


def main():
    parser = argparse.ArgumentParser(description='drgn cgroup tools')

    parser.add_argument('-C', '--cgroup_addr', dest='cgroup_addr', action='store', help='target cgrup sys state')
    parser.add_argument('-p', '--pid', dest='task_pid', action='store', help='target task pid')
    args = parser.parse_args()


    print()
    if args.cgroup_addr != None:
        print_cgroup_from_addr(args.cgroup_addr)
    elif args.task_pid != None:
        print_cgroup_from_pid(int(args.task_pid))
    else:
        task = per_cpu(prog['runqueues'], prog['crashing_cpu']).curr
        print_cgroup_task(task)


if __name__ == '__main__':
    main()
